// init.c 20220903 Renzis
#include <arch/cpu.h>
#include <arch/gdt.h>
#include <arch/interrupts.h>
#include <mm/pgtable.h>
#include <string.h>
#include <mm/mm.h>
#include <mm/mmio.h>
#include <sys/unistd.h>
#include <common.h>
#include <process.h>
#include <arch/boot/multiboot.h>
#include <keyboard.h>
#include <hal/pci.h>
#include <tty/console.h>
#include <tty/tty.h>
#include <mm/slab.h>
#include <mm/valloc.h>
#include <hal/apic.h>
#include <hal/ioapic.h>
#include <hal/ahci/ahci.h>
#include <hal/acpi/acpi.h>
#include <arch/vectors.h>
#include <timer.h>
#include <lib/math.h>
#include <fs/devfs.h>
#include <fs/ramfs.h>
#include <device/block.h>
#include <device/blkio.h>
#include <mm/vma.h>

#define LOG_TAG "init"
#include <log.h>

// Provided by linker (see linker.ld)
extern uptr __kernel_start;
extern uptr __kernel_mem_end;
extern uptr __init_hhk_start;
extern uptr __init_hhk_end;
extern uptr __user_start;
extern uptr __user_end;

multiboot_info_t* mb;

void print_multiboot_info() {
    printk("flags: %x\n", mb->flags);
    if (mb->flags & MULTIBOOT_INFO_MEMORY) {
        printk("mem_lower: %x\n", mb->mem_lower);
        printk("mem_upper: %x\n", mb->mem_upper);
    }
    if (mb->flags & MULTIBOOT_INFO_BOOTDEV)
        printk("boot_device: %x\n", mb->boot_device);
    // printk("cmdline: %x\n", mb->cmdline);
    if (mb->flags & MULTIBOOT_INFO_MODS) {
        printk("mods_count: %u\n", mb->mods_count);
        printk("mods_addr: %x\n", *(uptr*)mb->mods_addr);
    }
    // printk("mmap_length: %u\n", mb->mmap_length);
    // printk("mmap_addr: %x\n", *(uptr *)mb->mmap_addr);
    if (mb->flags & MULTIBOOT_INFO_DRIVE_INFO) {
        printk("drives_length: %u\n", mb->drives_length);
        printk("drives_addr: %x\n", *(uptr*)mb->drives_addr);
    } else {
        printk("No drives\n");
    }
    // printk("config_table: %x\n", mb->config_table);
    // printk("boot_loader_name: %s\n", (char *)mb->boot_loader_name);
    // printk("apm_table: %x\n", mb->apm_table);
    // printk("vbe_control_info: %x\n", mb->vbe_control_info);
    // printk("vbe_mode_info: %x\n", mb->vbe_mode_info);
    // printk("vbe_mode: %u\n", mb->vbe_mode);
    // printk("vbe_interface_seg: %u\n", mb->vbe_interface_seg);
    // printk("vbe_interface_off: %u\n", mb->vbe_interface_off);
    // printk("vbe_interface_len: %u\n", mb->vbe_interface_len);
    // printk("framebuffer_addr: %llu\n", mb->framebuffer_addr);
    // printk("framebuffer_pitch: %u\n", mb->framebuffer_pitch);
    // printk("framebuffer_width: %u\n", mb->framebuffer_width);
    // printk("framebuffer_height: %u\n", mb->framebuffer_height);
    // printk("framebuffer_bpp: %u\n", mb->framebuffer_bpp);
    // printk("framebuffer_type: %u\n", mb->framebuffer_type);

    // 根据选择的 framebuffer_type 打印额外字段
    // if (mb->framebuffer_type == MULTIBOOT_FRAMEBUFFER_TYPE_INDEXED) {
    //     printk("framebuffer_palette_addr: %x\n", mb->framebuffer_palette_addr);
    //     printk("framebuffer_palette_num_colors: %u\n", mb->framebuffer_palette_num_colors);
    // } else if (mb->framebuffer_type == MULTIBOOT_FRAMEBUFFER_TYPE_RGB) {
    //     printk("framebuffer_red_field_position: %u\n", mb->framebuffer_red_field_position);
    //     printk("framebuffer_red_mask_size: %u\n", mb->framebuffer_red_mask_size);
    //     printk("framebuffer_green_field_position: %u\n", mb->framebuffer_green_field_position);
    //     printk("framebuffer_green_mask_size: %u\n", mb->framebuffer_green_mask_size);
    //     printk("framebuffer_blue_field_position: %u\n", mb->framebuffer_blue_field_position);
    //     printk("framebuffer_blue_mask_size: %u\n", mb->framebuffer_blue_mask_size);
    // }
}

void mem_init() {
    page_init();
    mm_init(MEM_1MB + (mb->mem_upper << 10));

    for (multiboot_memory_map_t* mmap = (multiboot_memory_map_t*)mb->mmap_addr;
        (uptr)mmap < mb->mmap_addr + mb->mmap_length;
        mmap = (multiboot_memory_map_t*)((uptr)mmap + mmap->size + sizeof(mmap->size))
    ) {
        if (mmap->type == MULTIBOOT_MEMORY_AVAILABLE) {
            mark_chunk_free((mmap->addr_low + 0x0fffU) >> PG_SIZE_BITS, mmap->len_low >> PG_SIZE_BITS);
        }
    }

    size_t page_count = V2P(&__kernel_mem_end) >> PG_SIZE_BITS;
    mark_chunk_used(0, page_count, PG_LOCKED);

    mark_chunk_used(VGA_BUF_PA >> PG_SIZE_BITS, VGA_BUF_SIZE >> PG_SIZE_BITS, PG_LOCKED);
}

void reserve_memory(u32 free) {
    for (multiboot_memory_map_t* mmap = (multiboot_memory_map_t*)mb->mmap_addr;
        (uptr)mmap < mb->mmap_addr + mb->mmap_length;
        mmap = (multiboot_memory_map_t*)((uptr)mmap + mmap->size + sizeof(mmap->size))
    ) {
        uptr pa = PG_ALIGN(mmap->addr_low);
        if (mmap->type == MULTIBOOT_MEMORY_AVAILABLE || pa <= MEM_1MB) {
            continue;
        }
        size_t page_num = CEIL(mmap->len_low, PG_SIZE_BITS);
        size_t j = 0;
        if (free) {
            for (; j < page_num; j++) {
                uptr addr = pa + (j << PG_SIZE_BITS);
                unmapping_p2v(addr);
                if (mmap->type == MULTIBOOT_MEMORY_ACPI_RECLAIMABLE) {
                    mark_page_free(addr >> PG_SIZE_BITS);
                }
            }
        } else {
            for (; j < page_num; j++) {
                uptr addr = pa + (j << PG_SIZE_BITS);
                if (addr >= KERNEL_MM_BASE) {
                    break;
                }
                mapping_p2v(addr, addr, PG_PRESENT, VMAP_NULL);
                mark_page_used(addr >> PG_SIZE_BITS, PG_LOCKED);
            }
        }
    }
}

void kernel_init();
void kernel_main();
void assign_return_context(u32 return_to_kernel, struct process* p, uptr stk_top);
void build_process0() {
    uptr pt = get_page(PG_KERNEL);
    uptr stk = get_pages(2, PG_KERNEL);
    PDE_ENTRY(DIR_BASE_VADDR, KSTK_AREA_START >> 12) = (pg_entry_t){
        .frame = pt >> PG_SIZE_BITS,
        .present = 1,
        .rw = 1,
        .user = 0,
    };

    memset(PT_BASE(TBL_BASE_VADDR, KSTK_AREA_START >> 12), 0, PG_SIZE);

    PTE_ENTRY(TBL_BASE_VADDR, KSTK_AREA_END >> 12) = (pg_entry_t){
        .frame = (stk + 0x1000) >> PG_SIZE_BITS,
        .present = 1,
        .rw = 1,
        .user = 0,
    };
    PTE_ENTRY(TBL_BASE_VADDR, (KSTK_AREA_END - KSTK_SIZE + 0x10) >> 12) = (pg_entry_t){
        .frame = stk >> PG_SIZE_BITS,
        .present = 1,
        .rw = 1,
        .user = 0,
    };

    assign_return_context(1, cur_proc, KSTK_AREA_END);

    cur_proc->context.ctx->ctx.eip = (u32)kernel_init;
    cur_proc->context.ctx->ctx.eflags.flags.IF = 1;

    // 这里我们通过iret pop出eflags来开启中断
    asm volatile("pushl %0\n"
        "jmp switch_to\n" ::"r"(cur_proc)
        : "memory");
}

void reallocation_address_space() {
    pg_table_t* dir = (pg_table_t*)DIR_BASE_VADDR;
    uptr kernel_i = DIRECTORY_INDEX(KERNEL_MM_BASE);
    for (size_t i = kernel_i; i < KSTK_AREA_DIR_IDX; i++) {
        if (dir->entry[i].frame) {
            continue;
        }
        uptr pa = get_page(PG_KERNEL);
        assert_msg(pa, "Not enough page for kernel space");
        dir->entry[i] = (pg_entry_t){
            .frame = pa >> PG_SIZE_BITS,
            .present = 1,
            .rw = 1,
            .user = 0
        };
        memset(PT_BASE(TBL_BASE_VADDR, i << PG_INDEX_BITS), 0, PG_SIZE);
    }

    // Make 'NULL'
    pg_entry_t* pte = &PTE_ENTRY(TBL_BASE_VADDR, 0);
    mark_page_free(pte->frame);
    pte->data = 0;

    uptr hhk_init_start = (uptr)&__init_hhk_start;
    uptr hhs_init_end = (uptr)&__init_hhk_end;
    for (size_t i = hhk_init_start; i < hhs_init_end; i += PG_SIZE) {
        unmapping_p2v(i);
    }

    for (uptr i = (uptr)&__user_start; i < (uptr)&__user_end; i += PG_SIZE) {
        pte = &PDE_ENTRY(TBL_BASE_VADDR, i >> PG_SIZE_BITS);
        pte->user = 1;
        pte = &PTE_ENTRY(TBL_BASE_VADDR, i >> PG_SIZE_BITS);
        pte->user = 1;
        cpu_flush_page(i);
    }
}

void kernel_pre_init(multiboot_info_t* mb_info) {
    init_interrupts();

    mb = mb_info;
    mem_init();
    kmem_cache_init();
    valloc_init();

    tty_init((u16*)ioremap(VGA_BUF_PA, 1));
    tty_set_theme(VGA_COLOR_LIGHT_WHITE, VGA_COLOR_BLACK);
    console_init();

    fs_init();
    process_init();
    fs_mount_root("ramfs", NULL);
    fs_mount("/dev", "devfs", NULL, 0);
}

void kernel_init() {
    if (fork() == 0) {
        while (1) {
            // yield(); // idling!
        }
    }

    reserve_memory(0);

    acpi_init();
    apic_init();
    ioapic_init();
    pci_init();
    keyboard_init();
    timer_init();
    console_register();
    block_init();
    blkio_init();
    ahci_init();
    vma_init();

    reserve_memory(1);

    /**
     * 将所有内核空间的页目录添加相应页表映射
     * 这样可以保证任何进程(地址空间)在内核态的动作
     * 都会被同步到其他进程
     */
    reallocation_address_space();
    swapon("/dev/sdap2");

    mkdir("/root");
    mount("/dev/sdap1", "/root", "fat32", 0);
    
    kernel_main();
}